

import torch  # 张量可视化 TensorVisualization 显示的就是torch的张量
import matplotlib
from matplotlib import pyplot as plt

from matplotlib import colors
import numpy as np
import cv2 as cv # 借助opencv来显示曲线图片。


class Curve:  
    
    def __init__(self, xlabel="", ylabel="", legend=[],
                 xlim=(0,100),ylim=(0,100), xscale='linear', yscale='linear',
                 figsize=(3.5, 2.5)):
        """_summary_
           
        Args:
            xlabel (字符串): _description_. 描述x轴标题.
            ylabel (字符串): _description_. 描述y轴标题.
            legend (字符串列表): _description_. 描述各条曲线的标签（与样式fmts结合显示）.这个列表的长度决定了绘制几条曲线。
            xlim (元组类型): _description_. 定义坐标轴起止坐标（x_min,x_max）.
            ylim (元组类型): _description_. 定义坐标轴起止坐标（y_min,y_max）.
            xscale (str): _description_. Defaults to 'linear'.描述矩阵的坐标轴类型。
            yscale (str): _description_. Defaults to 'linear'.'linear'线性刻度；'log'对数刻度，方便显示[0,∞);'symlog'对称对数刻度，方便显示[-∞,∞)
            figsize (tuple): _description_. Defaults to (3.5, 2.5). 图片尺寸，单位是英寸.
        """
        # 曲线的绘图样式： 例如'g-'表示绿色实线。具体符号意义：'-'实线，'--'虚线，'-.'点画线，':'点线，'b'默认蓝，'g'绿，'r'红，'c'青，'m'洋红，'k'黑，'y'黄，'w'白
        plt.cla() 
        default_fmt = [] # 默认样式。
        line_color = ['b','g','r','c','m','k']
        line_style = [':','-.','--','-']
        for s in line_style:
            for c in line_color:
                default_fmt.append(c+s)
        #print(default_fmt)

        self.line_count= min( len(legend),len(default_fmt))   # 曲线数量不超过样式数量。
        self.points_x, self.points_y, self.fmts = [], [], []  # points_xy是列表的列表，[[曲线1],[曲线2]…… ]
        for i in range(self.line_count):
            self.points_x.append([])   
            self.points_y.append([])
            self.fmts.append(default_fmt[i])
        
        self.figsize = figsize
        self.xlabel, self.ylabel, self.xlim, self.ylim, self.xscale, self.yscale, self.legend = (
            xlabel, ylabel, xlim, ylim, xscale, yscale, legend)

    def add(self, x, y):
        """_summary_
        Args:
            x (标量，或数字列表):如果是标量，绘制多条曲线时，x坐标都一样。如果是列表，则有对应的y点才给对应的曲线添加点
            y (标量，或数字列表):如果是标量，只给第一条曲线添加点。如果是列表，则给对应的曲线添加点。
        """

        # 考虑x或y形参是标量的情况。
        if not hasattr(y, "__len__"): 
            y = [y]
        n = len(y)
        if not hasattr(x, "__len__"):
            x = [x] * n
        # 考虑x,y都是列表的情况，列表长度可能不一致。这时截取成短的列表。
        n = min(n, len(x),self.line_count)
        x = x[:n]
        y = y[:n]   
        # 现在x,y 是长度一样的列表。

        for i, (a, b) in enumerate(zip(x, y)):
            # enumerate函数用于在遍历可迭代对象时同时获取元素的索引和元素本身。
            # 当应用于zip(x, y)时，它会给由zip函数生成的每个元组添加一个索引。
            # 例如，对于前面zip的结果，enumerate(zip(x, y))会生成类似
            # [(0, (n1_x, n1_y)), (1, (n2_x, n2_y)), (2, (n3_x, n3_y))]的结果。
            if a is not None and b is not None:
                self.points_x[i].append(a) #将xy坐标添加到类的曲线坐标中。 X，Y都是二维数组。
                self.points_y[i].append(b)

        self.fig, self.axes = plt.subplots(1, 1, figsize=self.figsize)
        self.axes.cla()  # 清除当前坐标轴中的所有内容

        for x, y, fmt in zip(self.points_x, self.points_y, self.fmts):
            self.axes.plot(x, y, fmt)
        #self.config_axes()
        
        self.axes.set_xlabel(self.xlabel)
        self.axes.set_ylabel(self.ylabel)

        scale=self.xscale
        if scale == 'linear':  self.axes.set_xscale('linear')
        elif scale == 'log':   self.axes.set_xscale('log')
        elif scale == 'symlog':self.axes.set_xscale('symlog')
        scale=self.yscale
        if scale == 'linear':  self.axes.set_yscale('linear')
        elif scale == 'log':   self.axes.set_yscale('log')
        elif scale == 'symlog':self.axes.set_yscale('symlog')
               
        self.axes.set_xlim(self.xlim)
        self.axes.set_ylim(self.ylim)
        self.axes.legend(self.legend)
        self.axes.grid()

    def save(self,curve_img_name="curve_output.png"):
        self.fig.savefig(curve_img_name)
        #plt.pause(10)  #单位是秒

    def show_and_waitKey(self):
        """_summary_ 显示图片，等待键盘按键按下再退出显示
        """
        file_name="temp.png"
        self.save(file_name)

        img = cv.imread(file_name)
        # cv.imshow 第一个参数是窗口名，它是一个字符串,只能全英文字符，中文字符会出问题，第二个参数就是我们的图像。
        # 可以根据需要创建任意数量的窗口，但是窗口名字要不同。
        if img != None:
            cv.imshow(file_name, img)

        # cv.waitKey(0) 该函数为任意键盘事件等待指定毫秒。
        # 如果你在这段时间内按下任意键，程序将继续。
        # 如果传的是 0，它会一直等待键盘按下。
        cv.waitKey(0) # 等待按键按下再继续执行，毫秒为单位。返回按键的值。

        cv.destroyAllWindows() #销毁所有创建的窗口

#import seaborn as sns  # pip install seaborn

    """_summary_ 
    张量可视化，张量最好不要超过四维，张量维度太高影响直观性。
    这里干脆放弃4维度以上的张量可视化了。
    (图行，图列，行，列)，（图列，行，列），（行，列），（列）
    """

def TensorVisualization(tensor, fig_size=(5,5),image_file_name="tensor.png"):
    """_summary_ 将张量转换成图片并保存。
    张量可视化，张量最好不要超过四维，张量维度太高影响直观性。
    这里干脆放弃4维度以上的张量可视化了。
    (图行，图列，行，列)，（图列，行，列），（行，列），（列）
    Args:
        tensor (_type_): _description_ torch的张量类型。
        fig_size (tuple, optional): _description_. 生成图片的大小，单位是英寸 Defaults to (5,5).
        image_file_name (str, optional): _description_. 生成图片的保存路径 Defaults to "tensor.png".
    """

    plt.clf() # 清除当前图形
    plt.close('all') # 关闭所有窗口

    dim_len = len(tensor.shape)
    if dim_len > 4:
        print("张量维度太高，可视化失败") 
        return
    if dim_len < 1 :
        print("张量维度太低，可视化失败")
        return
    # 下面将维度统一成四维
    if dim_len == 3:
        tensor = tensor.reshape((1,tensor.shape[0],tensor.shape[1],tensor.shape[2]))
    if dim_len == 2:
        tensor = tensor.reshape((1,1,tensor.shape[0],tensor.shape[1]))
    if dim_len == 1:
        tensor = tensor.reshape((1,1,1,tensor.shape[0]))
    subfig_rows = tensor.shape[0] 
    subfig_cols = tensor.shape[1] # 最后两个维度以热力图的形式表现出来。 这个热力图又可以排列成二维矩阵，以此表现四维张量。
    # 归一化色条范围
    norm = colors.Normalize(vmin=tensor.min(), vmax=tensor.max())
    fig, axes = plt.subplots(subfig_rows, subfig_cols, figsize=fig_size)

    if subfig_rows * subfig_cols == 1:
        ax = [axes]
    else:
        ax = axes.flatten()  # type: ignore

    im =None
    for i in range(len(ax)):
        img = tensor[i//subfig_cols,i%subfig_cols,:,:].numpy()  #  ‘//’ 表示整除
        ax[i].xaxis.set_visible(False)
        ax[i].yaxis.set_visible(False) # X x轴设置不可见
        im = ax[i].imshow(img,norm=norm,cmap ='viridis')

    '''
    im = ax[i].imshow(img,norm=norm,cmap ='viridis') 的 cmap 参数解析

定性（Qualitative）颜色映射
Pastel1、Pastel2、Paired、Accent、Dark2、Set1、Set2、Set3：
这些颜色映射适用于区分不同的类别，每个类别会被分配一个独特的颜色。
例如，Set1有 9 种鲜明且容易区分的颜色，在绘制分类数据（如不同种类的动物、植物品种等）的图像或图表时非常有用。
如果要展示一个包含 9 个不同类别的柱状图，使用Set1颜色映射可以清晰地为每个类别分配一个独特的颜色，让观众能够轻松区分不同类别。

tab10、tab20、tab20b、tab20c：
这组颜色映射也是用于分类数据的可视化。tab10提供了 10 种颜色，tab20有 20 种颜色，而tab20b和tab20c是对tab20颜色的不同排列组合。
它们在绘制包含多个分类的图形（如饼图、堆叠柱状图等）时很有用，能够确保不同的分类有清晰可辨的颜色。


顺序（Sequential）颜色映射
Greys、Purples、Blues、Greens、Oranges、Reds、Yellows、Browns、Greys_r、Purples_r、Blues_r、Greens_r、Oranges_r、Reds_r、Yellows_r、Browns_r：
这些颜色映射呈现出从浅到深（没有_r后缀的）或从深到浅（有_r后缀，表示反向）的颜色顺序变化。
它们适用于表示具有顺序关系的数据，例如海拔高度、温度变化等。
例如，在绘制一个显示山脉海拔高度的地形图时，使用Greys颜色映射可以将低海拔区域用浅灰色表示，高海拔区域用深灰色表示，从而直观地展示海拔的高低变化。
viridis、plasma、inferno、magma、cividis、viridis_r、plasma_r、inferno_r、magma_r、cividis_r：
这是一组较新的、在感知上更均匀的顺序颜色映射。viridis颜色映射从深蓝色到黄色渐变，在科学可视化领域应用广泛，因为它在颜色过渡上更符合人眼的感知，即使对于色盲人群也相对友好。
例如，在展示一个海洋温度变化的热力图时，使用viridis颜色映射可以很好地呈现温度从低到高的变化。


发散（Diverging）颜色映射
PiYG、PRGn、BrBG、PuOr、RdGy、RdBu、RdYlBu、RdYlGn、Spectral、PiYG_r、PRGn_r、BrBG_r、PuOr_r、RdGy_r、RdBu_r、RdYlBu_r、RdYlGn_r、Spectral_r：
发散颜色映射的特点是有一个中心值（通常是中间色调），两侧的颜色向不同方向发散，用于表示数据在中心值两侧的变化情况，例如表示数据的正负偏差。
在绘制表示利润和亏损的热力图时，使用RdBu颜色映射可以将亏损部分用蓝色表示，盈利部分用红色表示，中间的平衡区域（如收支平衡）用白色或灰色表示，从而清晰地展示数据的正负差异。
    '''

    plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域
    if im is not None:
        fig.colorbar(im, ax=[ax[i] for i in range(subfig_rows*subfig_cols)])
    plt.savefig(image_file_name, bbox_inches='tight')
    #plt.show()

    plt.clf() # 清除当前图形


if __name__ == "__main__":

    """
    """
    c = Curve(xlabel= "x",ylabel="y",legend=["L1","L2"])
    c.add(20,[30,40,50])
    c.add(30,[40,50,])
    c.add(40,[90])
    c.add(50,[60,70,80])
    c.save()
    #c.show_and_waitKey()

    X = torch.rand(2,10)
    print(X)
    print(X[0])

    #test_heat_map2()

    t = torch.rand(3,3,67,45)*20-5
    #t[1,1,:,:]=torch.rand(23,45)-20
    #print(t[1,1,:,:])
    #print(t[2,2,:,:])
    file_name = "tensor.png"
    TensorVisualization(t,image_file_name=file_name)

    if 0 : # 在终端机上，没有图形界面，无法显示。下面的代码会报错。
        img = cv.imread(file_name)
        if img is not None:
            cv.imshow(file_name, img)
        cv.waitKey(0) # 等待按键按下再继续执行，毫秒为单位。返回按键的值。
        cv.destroyAllWindows() #销毁所有创建的窗口





